# -*- coding: utf-8 -*-
'''
Created on 28.08.2019

@author: yu03
'''


import datetime
import cv2
import matplotlib.pyplot as plt
from scipy import signal
from FFT_Interpolation import FFT_interpolation_2, line_cal, FFT_interpolation_boxcar
import numpy as np
import os
from scipy.optimize import curve_fit
from Video_Experiment import file_name, txt_name, video_name
from mpl_toolkits.mplot3d import Axes3D

def fit_func(x, a, b, c):
    return a*(x-b)**2 + c

now = datetime.datetime.now()

fs_cam = 50

Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0

img_set = []
hor_angle_centers = []
ver_angle_centers = []
hor_phase_centers = []
ver_phase_centers = []

ver_freq_set = []
hor_m_k_num_set = []
ver_m_k_num_set = []
hor_f_fit_set = []
ver_f_fit_set = []
hor_phase_centers = []
hor_freq_set = []


''' 
    Reading from Video file
'''
cap = cv2.VideoCapture(file_name)
frame_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print(frame_num, "Frames")
print('ok')
for k in range(frame_num):
    ret, frame = cap.read() 
    img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    img_set.append(img)
#     cv2.imshow('frame',img)
#     if cv2.waitKey(1) & 0xFF == ord('q'):
#         break
# cap.release()
# cv2.destroyAllWindows()

frame_num = 0
hor_threshold, ver_threshold = 100, 100
hor_centerline, ver_centerline = img_set[int(frame_num/2)][int(len(img_set[0])/2), :], img_set[int(frame_num/2)][:, int(len(img_set[0][1])/2)]
hor_start, ver_start = np.where(hor_centerline > hor_threshold)[0][0], np.where(ver_centerline > ver_threshold)[0][0]
hor_end, ver_end = np.where(hor_centerline > hor_threshold)[0][-1], np.where(ver_centerline > ver_threshold)[0][-1]
hor_index = int((ver_start+ver_end)/2)
ver_index = int((hor_start+hor_end)/2)

line_num = 3
scale_factor = 1
hor_lines = np.linspace(hor_index-scale_factor*(line_num-1)/2, hor_index+scale_factor*(line_num-1)/2, line_num, dtype='int')
# hor_lines = np.linspace(470, 500, 31, dtype='int')
ver_lines = np.linspace(ver_index-scale_factor*(line_num-1)/2, ver_index+scale_factor*(line_num-1)/2, line_num, dtype='int')
print(hor_lines)
print(ver_lines)

hor_f_set, ver_f_set = [], []
hor_phi_set, ver_phi_set = [], []

for img in img_set[100:110]:
# for img in img_set[0:2]:
    frame_num += 1
    for i in range(line_num):
        hor_center = np.array(img[hor_lines[i]])
        ver_center = img[:, ver_lines[i]]
        
        hor_f_fit, hor_phase_estim, hor_m_k_num = line_cal(hor_center)
        ver_f_fit, ver_phase_estim, ver_m_k_num = line_cal(ver_center)
        
        hor_f_set.append(hor_f_fit)
        hor_phi_set.append(hor_phase_estim)
        ver_f_set.append(ver_f_fit)
        ver_phi_set.append(ver_phase_estim)
        
        print(frame_num, hor_m_k_num, ver_m_k_num)
        
hor_f_set = np.array(hor_f_set)
ver_f_set = np.array(ver_f_set)
hor_f_set = hor_f_set.reshape(line_num, frame_num, order='F')
ver_f_set = ver_f_set.reshape(line_num, frame_num, order='F')
hor_angle_set = (V_x-Lamda*hor_f_set/2)*1e6 ### urad
ver_angle_set = (V_x-Lamda*ver_f_set/2)*1e6 ### urad

print(hor_angle_set)

# hor_non_f_set, hor_non_phi_set = [], []
# ver_non_f_set, ver_non_phi_set = [], []
# for i in range(line_num):
#     hor_non_f, hor_non_phi = FFT_interpolation_boxcar(hor_angle_set[i]-np.average(hor_angle_set[i]), 1/frame_num)[0:2]
#     ver_non_f, ver_non_phi = FFT_interpolation_boxcar(ver_angle_set[i]-np.average(ver_angle_set[i]), 1/frame_num)[0:2]
#     hor_non_f_set.append(hor_non_f)
#     hor_non_phi_set.append(hor_non_phi)
#     ver_non_f_set.append(ver_non_f)
#     ver_non_phi_set.append(ver_non_phi)
    
# hor_amp = []
# ver_amp = []
# for i in range(line_num):
#     hor_amp_line = np.max(hor_angle_set[i])-np.min(hor_angle_set[i])
#     ver_amp_line = np.max(ver_angle_set[i])-np.min(ver_angle_set[i])
#     hor_amp.append(hor_amp_line)
#     ver_amp.append(ver_amp_line)
# 
# hor_amp = np.array(hor_amp)
# ver_amp = np.array(ver_amp)
# hor_used_line = hor_lines[np.where(hor_amp < 3)[0]]
# ver_used_line = ver_lines[np.where(ver_amp < 3)[0]]
# print (repr(hor_used_line))
# print (repr(ver_used_line))

# plt.figure(1)
# for i in range(line_num):
#     plt.plot(hor_angle_set[i])
# plt.grid(which='both', axis='both')
# plt.show()

# plt.figure(2)
# im = plt.imshow(hor_angle_set, cmap='gray', extent=[0, frame_num, hor_lines[0], hor_lines[-1]])
# plt.colorbar(im, fraction=0.046, pad=0.04)



# fig = plt.figure(3)
# ax = Axes3D(fig) 
# X = hor_lines
# Y = np.arange(frame_num)
# X, Y = np.meshgrid(Y, X)
# Z = hor_angle_set 
# ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=plt.get_cmap('rainbow'))
# ax.contourf(X, Y, Z, zdir='x', offset=-2, cmap=plt.get_cmap('rainbow'))
# plt.draw()

# y = [np.arange(frame_num)]*line_num
# y = np.array(y)
#  
# fig = plt.figure('hor_tilt')
# x = [hor_lines]*frame_num
# x = np.array(x).T
# ax = fig.gca(projection='3d')
# plt.gca().patch.set_facecolor('white')
# for i in range(line_num):
#     ax.plot(x[i,:],y[i,:],hor_angle_set[i,:])
#  
# fig = plt.figure('ver_tilt')
# x = [ver_lines]*frame_num
# x = np.array(x).T
# ax = fig.gca(projection='3d')
# plt.gca().patch.set_facecolor('white')
# for i in range(line_num):
#     ax.plot(x[i,:],y[i,:],ver_angle_set[i,:])

# plt.figure('nonlinearity fft')
# plt.subplot(2,2,1)
# plt.plot(hor_lines, hor_non_f_set)
# plt.grid(which='both', axis='both')
# plt.subplot(2,2,3)
# plt.plot(hor_lines, (hor_non_phi_set))
# plt.grid(which='both', axis='both')
# plt.subplot(2,2,2)
# plt.plot(ver_lines, ver_non_f_set)
# plt.grid(which='both', axis='both')
# plt.subplot(2,2,4)
# plt.plot(ver_lines, (ver_non_phi_set))
# plt.grid(which='both', axis='both')
# 
# 
# plt.figure('nonlinearity amp')
# plt.subplot(2,1,1)
# plt.plot(hor_lines, hor_amp)
# plt.grid(which='both', axis='both')
# plt.subplot(2,1,2)
# plt.plot(ver_lines, ver_amp)
# plt.grid(which='both', axis='both')
# plt.show()